import socket
import struct
import crcmod
import time

class ModbusTCPClient:
    def __init__(self, host='localhost', port=10001, timeout=2):
        self.host = host
        self.port = port
        self.timeout = timeout
        self.socket = None

    def connect(self):
        try:
            self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.socket.settimeout(self.timeout)  # 设置超时时间
            self.socket.connect((self.host, self.port))
            return True
        except Exception as e:
            print(f"Connection error: {e}")
            return False

    def close(self):
        if self.socket:
            self.socket.close()

    def read_tcp_register(self, address, function_code=3, count=1, unit=1):
        # Build Modbus TCP request
        transaction_id = 1
        protocol_id = 0
        length = 6
        request = struct.pack(
            '>HHHBBHH',
            transaction_id,
            protocol_id,
            length,
            unit,
            function_code,
            address,
            count
        )

        try:
            # Send request
            self.socket.send(request)

            # Receive response
            response = self.socket.recv(1024)
            _, _, _, _, function_code, byte_count = struct.unpack('>HHHBB', response[:9])

            if function_code != 3:
                return None  # Error in response

            # Parse register values
            registers = struct.unpack(f'>{byte_count // 2}H', response[9:9 + byte_count])
            return registers
        except socket.timeout:
            print(f"Read timeout for address {address}")
            return None
        except Exception as e:
            print(f"Read error: {e}")
            return None
        
    def read_register(self, address, function_code=3, count=1, unit=1):
        # Build Modbus RTU Over TCP request
        request = struct.pack('>BBHH', unit, function_code, address, count)

        # Calculate CRC
        crc16 = crcmod.predefined.mkPredefinedCrcFun('modbus')
        crc = crc16(request)
        request += struct.pack('<H', crc)

        try:
            # Send request
            self.socket.send(request)

             # 下面是关键的收包逻辑
            buffer = bytearray()
            start_time = time.time()
            while True:
                # 非阻塞读取，每次最多读一批（可根据设备协议调整，推荐每次少量如32字节）
                chunk = self.socket.recv(2048)  # 使用 recv() 接收数据，缓冲区大小为 2048 字节
                if chunk:
                    buffer.extend(chunk)
                if len(buffer) >= 5:
                    # Verify CRC
                    response = buffer
                    received_crc = struct.unpack('<H', response[-2:])[0]
                    calculated_crc = crc16(response[:-2])
                    if received_crc == calculated_crc:
                        # Format response as HEX string
                        hex_response = " ".join(f"{byte:02X}" for byte in response)
                        return hex_response # 收到完整数据且CRC校验通过，直接返回结果
                # 超时判断
                if (time.time() - start_time) > self.timeout:
                    break

            # Receive response
            response = buffer
            if len(response) < 5:
                print(f"Read timeout for address {address}")
                return None

            # Verify CRC
            received_crc = struct.unpack('<H', response[-2:])[0]
            calculated_crc = crc16(response[:-2])
            if received_crc != calculated_crc:
                print(f"CRC error for address {address}")
                return None

            # Format response as HEX string
            hex_response = " ".join(f"{byte:02X}" for byte in response)
            return hex_response
        except socket.timeout:
            print(f"Read timeout for address {address}")
            return None
        except Exception as e:
            print(f"Read error: {e}")

    def scan_devices(self, start_address, end_address, function_code=3, unit=1):
        devices = []
        for address in range(start_address, end_address + 1):
            result = self.read_register(address, function_code, 1, unit)
            if result is not None:
                devices.append((address, result))
        return devices